import sys,os
from pylab import *

mp_vec = loadtxt('JLA_WMAP9_MP_vecs.txt')
sp_vec = loadtxt('JLA_WMAP9_SP_vecs.txt')
data_vec = loadtxt('JLA_WMAP9_vecs.txt')

def get_alphas_chains(chain_root, chain_num=6, first_chain=1, w0_idx=2, n_w=20, vecs=None, dataonly=False):
	if vecs is None:
		print('vecs are not set, exit!')
		sys.exit(0)

	if first_chain >= chain_num:
		print('check chain_num and first_chain!')
		sys.exit(0)

	for n in range(first_chain,chain_num+1):
		alphas = []
		chain_name = chain_root+'_'+str(n)+'.txt'
		print '==> processing chain: %s'%(chain_name)

		if os.path.isfile(chain_name):
			chain = loadtxt(chain_name)
			for i in range(chain.shape[0]):
				weight = chain[i,0]
				chi2 = chain[i,1]
				w = chain[i,w0_idx:(w0_idx+n_w)]
				temp = []
				temp.append(weight)
				temp.append(chi2)
				for i in range(n_w):
					temp.append(dot(vecs[i,:],1+w))
				alphas.append(temp)

			if dataonly:
				savetxt(chain_root+'_data_alpha_'+str(n)+'.txt',array(alphas),fmt='%10.5E')
			else:
				savetxt(chain_root+'_alpha_'+str(n)+'.txt',array(alphas),fmt='%10.5E')
		else:
			print 'no %s'%(chain_name)


for i in range(1,101):
	chain_root = '../'
#	os.system('cd '+chain_root+'; tar zxvf chain_'+str(i)+'.tar.gz'+';cd -')
	get_alphas_chains(chain_root+'/chain_'+str(i)+'/EoS_SN_CMB',vecs=sp_vec,chain_num=4,w0_idx=6,n_w=20,first_chain=1)
	#get_alphas_chains(chain_root+'/chain_'+str(i)+'/EoS_SN_CMB',vecs=data_vec,chain_num=4,w0_idx=6,n_w=20,first_chain=1)
